package alg

import (
	"crypto/ecdsa"
	"crypto/rand"
	"encoding/pem"
	"errors"
	"github.com/emmansun/gmsm/pkcs8"
	"github.com/emmansun/gmsm/sm2"
	"github.com/emmansun/gmsm/smx509"
	"math/big"
	"os"
)

type SM2KeyContext struct {
	Key *sm2.PrivateKey
}

var SM2_UID []byte = []byte{
	0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38,
	0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38,
}

type Option func(ctx *SM2KeyContext) error

// WithSM2PubKey 公钥
func WithSM2PubKey(key []byte) Option {
	return func(ctx *SM2KeyContext) error {
		if len(key) != 64 {
			return errors.New("error length of public key")
		}
		ctx.Key.Curve = sm2.P256()
		ctx.Key.X = new(big.Int).SetBytes(key[:32])
		ctx.Key.Y = new(big.Int).SetBytes(key[32:])
		return nil
	}
}

// WithSM2PrivateKey 私钥
func WithSM2PrivateKey(key []byte) Option {
	return func(ctx *SM2KeyContext) error {
		if len(key) != 32 {
			return errors.New("error length of private key")
		}
		ctx.Key.Curve = sm2.P256()
		ctx.Key.D = new(big.Int).SetBytes(key)
		return nil
	}
}

// WithSM2PEMCer X509证书
func WithSM2PEMCer(c []byte) Option {
	return func(ctx *SM2KeyContext) error {

		block, _ := pem.Decode(c)

		cer, err := smx509.ParseCertificate(block.Bytes)
		if err != nil {
			return err
		}

		key, ok := cer.PublicKey.(*ecdsa.PublicKey)
		if !ok {
			return errors.New("key block is not of type SM2")
		}

		ctx.Key.PublicKey = *key

		return nil
	}
}

// WithSM2PEMCertFile X509证书
func WithSM2PEMCertFile(file string) Option {
	return func(ctx *SM2KeyContext) error {

		buf, err := os.ReadFile(file)
		if err != nil {
			return err
		}

		block, _ := pem.Decode(buf)

		cer, err := smx509.ParseCertificate(block.Bytes)
		if err != nil {
			return err
		}

		key, ok := cer.PublicKey.(*ecdsa.PublicKey)
		if !ok {
			return errors.New("key block is not of type SM2")
		}

		ctx.Key.PublicKey = *key

		return nil
	}
}

// WithSM2X509Key X509文件私钥
func WithSM2X509Key(file string) Option {
	return func(ctx *SM2KeyContext) error {

		buf, err := os.ReadFile(file)
		if err != nil {
			return err
		}

		block, _ := pem.Decode(buf)

		pk, err := pkcs8.ParsePKCS8PrivateKeySM2(block.Bytes, nil)
		if err != nil {
			return err
		}

		ctx.Key = pk

		return nil
	}
}

func NewSM2KeyContext(options ...Option) (*SM2KeyContext, error) {
	var err error
	ctx := &SM2KeyContext{
		Key: new(sm2.PrivateKey),
	}
	for _, option := range options {
		err = option(ctx)
		if err != nil {
			break
		}
	}
	return ctx, err
}

func (ctx *SM2KeyContext) Sign(data []byte) ([]byte, error) {
	return ctx.Key.Sign(rand.Reader, data, sm2.NewSM2SignerOption(true, SM2_UID))
}

func (ctx *SM2KeyContext) Verify(data, sign []byte) bool {
	return sm2.VerifyASN1WithSM2(&ctx.Key.PublicKey, SM2_UID, data, sign)
}

func (ctx *SM2KeyContext) Encrypt(data []byte) ([]byte, error) {
	return sm2.Encrypt(rand.Reader, &ctx.Key.PublicKey, data, sm2.NewPlainEncrypterOpts(sm2.MarshalUncompressed, sm2.C1C3C2))
}

func (ctx *SM2KeyContext) Decrypt(data []byte) ([]byte, error) {
	return ctx.Key.Decrypt(rand.Reader, data, sm2.NewPlainDecrypterOpts(sm2.C1C3C2))
}

func GenSM2KeyPair() ([]byte, []byte, error) {
	keys, err := sm2.GenerateKey(rand.Reader)
	if err != nil {
		return nil, nil, err
	}
	return append(keys.X.Bytes(), keys.Y.Bytes()...), keys.D.Bytes(), nil
}
